import struct
import socket
import time
import random
from scapy.all import *
# import pdb

def generate_radius_packet(code=1, identifier=None, authenticator=None, attributes=None):
    """
    生成一个 RADIUS 数据包

    Args:
        code: RADIUS 包类型 (1=Access-Request, 2=Access-Accept, 3=Access-Reject)
        identifier: 包标识符
        authenticator: 验证器 (16字节)
        attributes: RADIUS 属性列表
    """
    # pdb.set_trace()
    if identifier is None:
        identifier = random.randint(0, 255)
    # pdb.set_trace()

    if authenticator is None:
        authenticator = bytes([random.randint(0, 255) for _ in range(16)])

    if attributes is None:
        attributes = []

    # 添加一些示例属性
    # User-Name 属性 (类型1)
    username = "testuser"
    attributes.append(struct.pack("!BB", 1, len(username) + 2) + username.encode())

    # User-Password 属性 (类型2)
    password = "testpassword"
    # 简化处理，实际中需要加密
    attributes.append(struct.pack("!BB", 2, len(password) + 2) + password.encode())

    # NAS-IP-Address 属性 (类型4)
    nas_ip = socket.inet_aton("192.168.1.1")
    attributes.append(struct.pack("!BB", 4, 6) + nas_ip)

    # NAS-Identifier 属性 (类型32)
    nas_id = "nas01.example.com"
    attributes.append(struct.pack("!BB", 32, len(nas_id) + 2) + nas_id.encode())

    # 连接所有属性
    attr_data = b"".join(attributes)
    # 计算总长度
    length = 20 + len(attr_data)  # RADIUS 头部长度20字节

    # 构造 RADIUS 头部
    radius_header = struct.pack("!BBH", code, identifier, length) + authenticator

    return radius_header + attr_data


def create_radius_traffic_pcap(filename="radius_traffic.pcap", packet_count=10):
    """
    生成 RADIUS 流量并保存为 pcap 文件

    Args:
        filename: 输出的 pcap 文件名
        packet_count: 生成的数据包数量
    """
    packets = []

    for i in range(packet_count):
        # 随机选择包类型
        code = random.choice([1, 2, 3])  # Access-Request, Access-Accept, Access-Reject

        # 生成 RADIUS 数据包
        radius_data = generate_radius_packet(code=code, identifier=i)

        # 创建完整数据包 (以太网帧)
        packet = Ether(dst="00:11:22:33:44:55", src="aa:bb:cc:dd:ee:ff") / \
                 IP(src="192.168.1.100", dst="192.168.1.1") / \
                 UDP(sport=1812, dport=1812) / \
                 Raw(radius_data)

        # 设置时间戳
        packet.time = time.time() + i

        packets.append(packet)

        # 模拟一些 Access-Request/Access-Accept 对
        if code == 1 and i < packet_count - 1:  # Access-Request
            # 对应的 Access-Accept 包
            accept_data = generate_radius_packet(code=2, identifier=i)
            accept_packet = Ether(dst="aa:bb:cc:dd:ee:ff", src="00:11:22:33:44:55") / \
                            IP(src="192.168.1.1", dst="192.168.1.100") / \
                            UDP(sport=1812, dport=1812) / \
                            Raw(accept_data)
            accept_packet.time = packet.time + 0.1
            packets.append(accept_packet)

    # 保存到 pcap 文件
    wrpcap(filename, packets)
    print(f"生成了 {len(packets)} 个 RADIUS 数据包到 {filename}")


def create_radius_accounting_traffic(filename="radius_accounting.pcap", packet_count=5):
    """
    生成 RADIUS Accounting 流量 (类型 4=Accounting-Request, 5=Accounting-Response)
    """
    packets = []

    for i in range(packet_count):
        # Accounting-Request (Start)
        acct_attrs = []

        # User-Name
        username = f"user{i:03d}"
        acct_attrs.append(struct.pack("!BB", 1, len(username) + 2) + username.encode())

        # NAS-Identifier
        nas_id = "nas01.example.com"
        acct_attrs.append(struct.pack("!BB", 32, len(nas_id) + 2) + nas_id.encode())

        # Acct-Status-Type (Start = 1)
        acct_attrs.append(struct.pack("!BBH", 40, 4, 1))

        # Acct-Session-Id
        session_id = f"session{i:05d}"
        acct_attrs.append(struct.pack("!BB", 44, len(session_id) + 2) + session_id.encode())

        # Acct-Authentic (RADIUS = 1)
        acct_attrs.append(struct.pack("!BBB", 45, 3, 1))

        # 连接属性
        attr_data = b"".join(acct_attrs)
        length = 20 + len(attr_data)
        authenticator = bytes([random.randint(0, 255) for _ in range(16)])
        radius_header = struct.pack("!BBH", 4, i * 2, length) + authenticator  # 4 = Accounting-Request
        radius_data = radius_header + attr_data

        # 创建数据包
        packet = Ether(dst="00:11:22:33:44:55", src="aa:bb:cc:dd:ee:ff") / \
                 IP(src="192.168.1.100", dst="192.168.1.1") / \
                 UDP(sport=1813, dport=1813) / \
                 Raw(radius_data)

        packet.time = time.time() + i * 2
        packets.append(packet)

        # 对应的 Accounting-Response
        resp_attrs = []
        resp_attrs.append(struct.pack("!BBH", 40, 4, 2))  # Acct-Status-Type (Stop = 2)
        resp_attr_data = b"".join(resp_attrs)
        resp_length = 20 + len(resp_attr_data)
        resp_radius_header = struct.pack("!BBH", 5, i * 2, resp_length) + authenticator  # 5 = Accounting-Response
        resp_radius_data = resp_radius_header + resp_attr_data

        resp_packet = Ether(dst="aa:bb:cc:dd:ee:ff", src="00:11:22:33:44:55") / \
                      IP(src="192.168.1.1", dst="192.168.1.100") / \
                      UDP(sport=1813, dport=1813) / \
                      Raw(resp_radius_data)

        resp_packet.time = packet.time + 0.1
        packets.append(resp_packet)

    # 保存到 pcap 文件
    wrpcap(filename, packets)
    print(f"生成了 {len(packets)} 个 RADIUS Accounting 数据包到 {filename}")


if __name__ == "__main__":
    # 生成标准 RADIUS 认证流量
    create_radius_traffic_pcap("radius_auth.pcap", 20)

    # 生成 RADIUS 计费流量
    create_radius_accounting_traffic("radius_acct.pcap", 10)

    print("RADIUS 流量 pcap 文件生成完成!")
